#include "kernel_operator.h"
#include "lib/matmul_intf.h"

using namespace AscendC;

class NLLLossSum {
 public:
  	__aicore__ inline NLLLossSum() {}
  	__aicore__ inline void Init(GM_ADDR x, GM_ADDR target, GM_ADDR weight,
                              GM_ADDR y ,
                              int32_t ignoreIndex, 
                              uint32_t coreNum, uint32_t sizeC,
                              uint32_t coreRowNum, uint32_t coreRowMod, TPipe* pipeIn) {
	this->pipe = pipeIn;

    this->ignoreIndex = ignoreIndex;
	this->sizeC = sizeC;
	this->coreNum = coreNum;

	this->coreRowNum = coreRowNum;
	this->coreRowMod = coreRowMod;

	this->coreIndex = GetBlockIdx();
	this->beginRow = coreRowNum * this->coreIndex;
	this->beginRow += (coreIndex < coreRowMod) ? coreIndex : coreRowMod;
	this->coreRow = (coreIndex < coreRowMod) ? coreRowNum+1 : coreRowNum;


	yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, sizeof(DTYPE_Y));
	pipe->InitBuffer(ReduceXBuf, 32);
	auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
	Duplicate(reduceXLocal,(float)0.0, 8);
	AscendC::SetAtomicNone();
	if (this->coreIndex==coreNum-1)
	{
		
	// 	printf("yGm before: %f\n",yGm.GetValue(0));
		DataCopy(yGm, reduceXLocal, 8);
	// 	printf("yGm after: %f\n",yGm.GetValue(0));
	// 	// PipeBarrier<PIPE_ALL>();
	}
	
	// DataCopy(yGm, reduceXLocal, 8);
	// PipeBarrier<PIPE_ALL>();


	// InitGlobalMemory(yGm, 8, (float)0.0);
	xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->beginRow * this->sizeC,
                        this->coreRow * this->sizeC * sizeof(DTYPE_X));
    weightGm.SetGlobalBuffer((__gm__ DTYPE_WEIGHT*)weight,
                             sizeC * sizeof(DTYPE_WEIGHT));
	targetGm.SetGlobalBuffer(
        (__gm__ DTYPE_TARGET*)target + this->beginRow,
        this->coreRow * sizeof(DTYPE_TARGET));

    this->roundSizeC = ((this->sizeC+7) / 8) * 8;
    this->roundCoreRow = ((this->coreRow+7) / 8) * 8;

    pipe->InitBuffer(AllWeightBuf, roundSizeC * sizeof(DTYPE_WEIGHT));

	pipe->InitBuffer(AllTargetBuf, roundCoreRow * sizeof(DTYPE_TARGET));
	// pipe.InitBuffer(AllXBuf, roundCoreRow * sizeof(DTYPE_TARGET));

	// pipe.InitBuffer(inQueueTarget, BUFFER_NUM, ITER_N * sizeof(DTYPE_TARGET));

	pipe->InitBuffer(GatherWeightBuf, roundCoreRow * sizeof(DTYPE_WEIGHT));
	// this->gatherWeight = GatherWeightBuf.Get<DTYPE_WEIGHT>();

	pipe->InitBuffer(InitIndexBuf, roundCoreRow * sizeof(DTYPE_TARGET));
	pipe->InitBuffer(InitIndexXBuf, roundCoreRow * sizeof(DTYPE_TARGET));

	// pipe.InitBuffer(IndexBuf, this->coreRow * sizeof(DTYPE_TARGET));

	

	pipe->InitBuffer(GatherXBuf, roundCoreRow * sizeof(DTYPE_X));
	pipe->InitBuffer(CopyInXBuf, roundCoreRow * sizeof(DTYPE_X) * 8);

	
	pipe->InitBuffer(WorkBuf, (roundCoreRow ) * sizeof(DTYPE_X));


	}

	__aicore__ inline void process() {

		// printf("----sum class----\n coreNum=%d, coreIdx=%d, coreRow=%d \n",this->coreNum,this->coreIndex,this->coreRow);
		// PipeBarrier<PIPE_ALL>();
		auto allTarget = AllTargetBuf.Get<DTYPE_TARGET>();
		auto allWeight =AllWeightBuf.Get<DTYPE_WEIGHT>();
		auto copyInX = CopyInXBuf.Get<DTYPE_X>();
		auto gatherX = GatherXBuf.Get<DTYPE_X>();
		auto indexLocal = InitIndexBuf.Get<DTYPE_TARGET>();
		auto indexLocalUint = indexLocal.ReinterpretCast<uint32_t>();
		auto indexXLocal = InitIndexXBuf.Get<DTYPE_TARGET>();
		auto gatherWeight = GatherWeightBuf.Get<DTYPE_WEIGHT>();
		auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
		auto workLocal = WorkBuf.Get<DTYPE_X>();
		AscendC::DataCopyPadParams padParams{false, 0, 0, 0};
        AscendC::DataCopyParams copyParams{1, 4, 0, 0};
		AscendC::DataCopyExtParams extParams{1, 4, 0, 0, 0};

		
		CreateVecIndex(indexLocal, (int32_t)0, coreRow);
		Muls(indexXLocal, indexLocal, (int32_t)sizeC , coreRow);
		Muls(indexLocal, indexLocal, 32 , coreRow);

		/*
			init yGm
		*/


		/*
			copy in target
		*/
		
		DataCopy(allTarget, targetGm, roundCoreRow);
		// PipeBarrier<PIPE_ALL>();

		/*
			copy in weight
		*/
		
		DataCopy(allWeight, weightGm, roundSizeC);
		// PipeBarrier<PIPE_ALL>();

		/*
			gather x
		*/
		PipeBarrier<PIPE_ALL>();
		Add(indexXLocal, indexXLocal, allTarget, coreRow);
		// PipeBarrier<PIPE_ALL>();
		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("[i,indexLocal,indexXLocal]= %d %d\n", i,indexLocal.GetValue(i),indexXLocal.GetValue(i));
		// }
        for (int i = 0; i < coreRow; i++) {
            AscendC::DataCopyPad(copyInX[i*8], xGm[indexXLocal.GetValue(i)], copyParams, padParams);
        }


		PipeBarrier<PIPE_ALL>();
		Gather(gatherX, copyInX, indexLocalUint, (uint32_t)0, this->coreRow);

		/*
			gather weight
		*/

		Muls(indexLocal, allTarget, 4 , coreRow);
		// PipeBarrier<PIPE_ALL>();
		Gather(gatherWeight, allWeight, indexLocalUint, (uint32_t)0, this->coreRow);
		// PipeBarrier<PIPE_ALL>();

		/*
			Compute
		*/

		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("[i,target,weight,x]= %d %d %f %f\n", i,allTarget.GetValue(i),gatherWeight.GetValue(i), gatherX.GetValue(i));
		// }
		Mul(gatherX, gatherX, gatherWeight, this->coreRow);

		ReduceSum(reduceXLocal, gatherX, workLocal, this->coreRow);
    	// printf("reducexLocal: %f \n",reduceXLocal.GetValue(0));
		float p = -1;
		Muls(reduceXLocal, reduceXLocal, p, this->coreRow);

		// PipeBarrier<PIPE_ALL>();

		/*
			CopyToWorkSpace
		*/
	 	AscendC::SetAtomicAdd<float>();
   		AscendC::DataCopy(yGm, reduceXLocal, 8);
		// AscendC::SetAtomicNone();

	}


private:
	uint32_t sizeC;
	int32_t ignoreIndex;
	uint32_t coreRowNum,coreRowMod,coreNum,roundSizeC,roundCoreRow;
	uint32_t beginRow,coreIndex,coreRow;

	AscendC::TPipe* pipe;
	AscendC::GlobalTensor<DTYPE_X> xGm;
	AscendC::GlobalTensor<DTYPE_TARGET> targetGm;
	AscendC::GlobalTensor<DTYPE_WEIGHT> weightGm;
	AscendC::GlobalTensor<DTYPE_Y> yGm;

	TBuf<QuePosition::VECCALC>  AllTargetBuf, AllWeightBuf;
	TBuf<QuePosition::VECCALC> InitIndexBuf,InitIndexXBuf,IndexBuf,GatherWeightBuf,GatherXBuf, CopyInXBuf;
	TBuf<QuePosition::VECCALC> ReduceXBuf,WorkBuf;
	TBuf<QuePosition::VECCALC> FinalXBuf;

	LocalTensor<DTYPE_WEIGHT> allWeight,gatherWeight;
	LocalTensor<DTYPE_TARGET> allTarget;
	LocalTensor<DTYPE_TARGET> initIndex;
	LocalTensor<DTYPE_X> gatherX;
	
};